import torch
import sounddevice as sd
import numpy as np
import random
import time

from artear2_3.dataset import ParamsDataset, get_files
from artear2_3.network import VAE

from artear2_3.simpledaw import SimpleDAW
from artear2_3.strophie import DAW


if __name__ == '__main__':
    
    z_dim = 100
    
    cwd = 'C:/dev_spa/DMuse/202202b3'
    
    config = {
            'dataset_config': {
                'train_sample_dir': cwd,
                'test_sample_dir': cwd,
                'train_label_dir': cwd,
                'test_label_dir': cwd,
                'compress_rate': 30.,
                'n_spaces': 30,
            },

            'validation_config': {
                'source_sample_rate': 44100.,
                'bpm': 120,
            },
            'batch_size': 1,
            'shuffle': True,
        }

    
    dataset = ParamsDataset(mode='test', **config['dataset_config'])
    G = VAE(input_dim=dataset.n_params, z_dim=z_dim)

    G.load_state_dict(torch.load('vae.pth', map_location=torch.device('cpu')))

    daw = DAW(backend='reaper', track_length=4, apply_immediately=True)
    daw.add_track()
    daw.add_fx(-1, 'Sylenth1 (LennarDigital)', plugin_path='C:/VST/64bit/Sylenth1.dll')
    
    simple_daw = SimpleDAW('C:/VST/64bit/Sylenth1.dll', sample_rate=config['validation_config']['source_sample_rate'])
    
    midi_dir = r"B:\muse_repo\MIDI"
    mid_files = get_files(midi_dir, 'mid')
    mid = random.choice(mid_files)
    simple_daw.load_midi(mid)
    daw.load_midi(mid, track_id=-1)


    with open('params_generated.txt', 'wt') as file:
        for i in range(2000):
            # params = list(G(torch.randn(1, z_dimension)).detach().numpy()[0, :])
            z = torch.randn(1, z_dim)
            params = list(G.decode(z).view(-1, dataset.n_params).detach().numpy()[0, :])
            file.write(','.join(str(param) for param in params) + '\n')
    
    for i in range(2000):
        # params = list(G(torch.randn(1, z_dimension)).detach().numpy()[0, :])
        
        z = torch.randn(1, z_dim)
        params = list(G.decode(z).view(-1, dataset.n_params).detach().numpy()[0, :])
        
        for i in range(len(params)):
            if params[i] >= 0.99:
                params[i] = 1.
            elif params[i] < 0.01:
                params[i] = 0.
        print(params)
        
        print('正在播放预测参数')
        daw.set_fx_params(params, -1, -1)
        time.sleep(2)
        daw.backend.play()
         
        time.sleep(5)
        
        '''simple_daw.set_params(params)
        audio = simple_daw.render(8.)
        sd.play(audio, config['validation_config']['source_sample_rate'], blocking=True)'''
        print('播放结束')

        

